import torch
import matplotlib
import matplotlib.pyplot as plt

# matplotlib.use()函数用于指定matplotlib使用的图形后端
# 'TkAgg'表示使用Tkinter作为图形界面工具包，Agg作为渲染引擎
# 该设置必须在导入matplotlib.pyplot之前执行才有效
# 主要用于确保图形能够正确显示在Tkinter窗口中
matplotlib.use('TkAgg')

x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])


class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred


# 可调用对象
model = LinearModel()

# 损失函数
# reduction="sum"：计算所有样本损失的总和
# 对于批次中的每个样本都会计算损失值，然后将这些损失值相加得到一个标量结果
criterion = torch.nn.MSELoss(reduction="sum")
# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

epochs = 500
train_loss = []
for epoch in range(500):
    # 计算预测值
    y_pred = model(x_data)
    # 计算损失
    loss = criterion(y_pred, y_data)
    train_loss.append(loss.item())
    # 梯度清零
    optimizer.zero_grad()
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()

print("w = ", model.linear.weight.item())
print("b = ", model.linear.bias.item())

x_test = torch.Tensor([[4.0]])
y_test = model(x_test)
print("y_pred = ", y_test.data)

plt.figure(figsize=(8, 8))
plt.plot(range(epochs), train_loss, label='Training Loss')
# plt.legend() 创建图例，自动获取绘图时设置的标签
# loc='upper right' 参数指定图例在图表中的位置为右上角
plt.legend(loc='upper right')
plt.show()
